import numpy as np
import oars
from oars.matrices import *
from oars.utils import *
from oars.pep import *
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
# matplotlib.use('nbagg')
import seaborn as sns
sns.set()
from matplotlib.colors import CenteredNorm
import cvxpy as cvx
from collections import defaultdict
!mkdir -p figs

# Data for matrices in cluster example (6x6)
n = 6
Z_fixed = {(r,c):0 for r in range(4,6) for c in range(3)}
Z_fixed[(3,1)] = 0
Z_fixed[(3,2)] = 0

Z, W = getSimilar(n, fixed_Z=Z_fixed, fixed_W=Z_fixed)

# Figure 3
t = 0.016*np.ones(n)
l = 0.00025*np.ones((n, n))
l[0, 3] = 0.01
l[3, 0] = 0.01  
fig = getGantt(t, l, Z, W, title="Cluster Algorithm Parallelism", itrs=5)
fig.show()
fig.write_image("figs/fig_gantt.pdf")

# Figure 4 helper
def cplot(W,Z):
    vmin = np.floor(np.min(np.minimum(W, Z)))
    cm = plt.cm.coolwarm
    fig, ax = plt.subplots(1, 2, figsize=(12, 6))
    x = -np.arange(n)
    y = np.arange(n)
    ax[0].pcolormesh(x, y, W, norm=CenteredNorm(), cmap=cm) #, vmin=vmin, vmax=2)
    ax[0].set_title(r"$W$")
    ax[0].axis('off')
    ax[0].set_aspect('equal')

    pc = ax[1].pcolormesh(x, y, Z, norm=CenteredNorm(), cmap=cm)
    ax[1].set_title(r"$Z$")
    ax[1].axis('off')
    ax[1].set_aspect('equal')

    # Add colorbar
    plt.tight_layout()
    
    fig.colorbar(pc, ax=ax.ravel().tolist())
    return fig

# Figure 4
n = 10
m = 2
Z, W = getBlockMin(n, m, objective=getMaxConnectivity)

fig = cplot(W, Z)

# Save as pdf

fig.savefig("figs/fig_block_fiedler.pdf")

# Figure 4
Z, W = getBlockMin(n, m, objective=getMinSLEM)

fig = cplot(W, Z)
fig.savefig("figs/fig_block_slem.pdf")

# Figure 4
Z, W = getBlockMin(n, m, objective=getMinResist)

fig = cplot(W, Z)
fig.savefig("figs/fig_block_resist.pdf")

# Figure 4
Z, W = getBlockMin(n, m, objective=getSimilar)

fig = cplot(W, Z)
fig.savefig("figs/fig_block_lw.pdf")

# Figure 5 (hangs if executed in Windows using Kaleido 0.1.0post1)
# try pip install kaleido==0.1.0post1
n = 6
t = np.ones(n)
l = np.ones((n,n))/10
t[0] = 6
t[2] = 6
t[3] = 6

for i in range(1, n//2+1):
    Zb, Wb = getBlockMin(n, i, objective=getMinResist)
    if i == 1:
        itrs = 6
    else:
        itrs = 7
    fig = getGantt(t, l, Zb, Wb, title=str(n//i)+"-Block", itrs=itrs)
    
    fig.write_image("figs/fig_misdp"+str(n//i)+"block.pdf")
    fig.show()

Z, W = getMinCycle(n, t=t, l=l, objective=getMinResist)

fig = getGantt(t, l, Z, W, title="Minimum Iteration", itrs=12)

fig.write_image("figs/fig_misdp_min_cycle.pdf")
fig.show()

# Figure 6 - contraction rate results
# fig_smoothstronglyconvex.ipynb
# Function to generate list of all test matrices
def getTestMatrices(n):
    cases = []
    Ws = []
    names = ['Full', 'MT']
    Lf, Wf = getFull(n)
    Mf = getMfromWCholesky(Wf)
    cases.append((Lf, Mf))
    Ws.append(Wf)
    Lmt, Wmt = getMT(n)
    Mmt = getMfromWCholesky(Wmt)
    cases.append((Lmt, Mmt))
    Ws.append(Wmt)
    i = n//2
    for obj in [getMinSLEM, getMinResist, getMaxConnectivity]:
        L, W = getBlockMin(n, i, objective=obj)
        M = getMfromWCholesky(W)
        cases.append((L, M))
        Ws.append(W)
        title = f"2-Block "+ obj.__name__
        names.append(title)
        
    return cases, names, Ws
    
def getMatrixContraction(cases, names, Lc=2, mu=1, gamma=0.5):
    '''Get contraction factors for matrix cases

    Args:
        cases (list): list of cases with L and M
        names (list): list of names
        Lc (float): Lipschitz constant
        mu (float): strong convexity parameter
        gamma (float): step size

    Returns:
        taus (dict): dictionary of contraction factors
    '''

    taus = {}
    for i, (L, M) in enumerate(cases):
        n = L.shape[0]
        
        ls = np.ones(n)*Lc
        mus = np.ones(n)*mu
        
        tau, _ = getReducedContractionFactor(L, M, ls, mus)
        taus[names[i]] = tau
        print(names[i], tau)
    return taus

# Figure 6 helper
def testFixedRate(start=4, end=31, lc=2, mc=1):
    tau_results = defaultdict(list) 
    for n in range(start, end, 2):
        print(f"Testing optimal rate on {n} nodes")
        ls = np.ones(n) * lc
        mus = np.ones(n) * mc
        cases, names, Ws = getTestMatrices(n)
        for i, (L, M) in enumerate(cases):
            tau = getReducedContractionFactor(L, M, ls, mus)
            tau_results[names[i]].append(tau)


    # Convert results to df
    tau_df = pd.DataFrame(tau_results)
    tau_df.index = range(start, end, 2)

    return tau_df

# Figure 6a
tau_df = testFixedRate(start=6, end=31)

Testing optimal rate on 6 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 8 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 10 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 12 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 14 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 16 nodes

# Plot figure 6a
fig, ax = plt.subplots()
tau_df.plot(ax=ax)

ax.set_xlabel("n")
ax.set_ylabel("Contraction")
ax.set_title(r"Fixed step size contraction with $n$ Lipschitz strongly monotone operators")
ax.legend(['Full', 'MT', '2-Block Min SLEM', r'2-Block Min Resistance (also min $\|Z-W\|$)', '2-Block Max Fiedler'], loc=(0.25,0.4))
fig.savefig("figs/fig6a.pdf", format="pdf", bbox_inches="tight")

# Figure 6b helper
def testFixedRateProj(start=4, end=31, lc=2, mc=1):
    tau_results = defaultdict(list) 
    for n in range(start, end, 2):
        print(f"Testing optimal rate on {n} nodes")
        ls = np.ones(n) * lc
        ls[n-1] = np.inf
        mus = np.ones(n) * mc
        mus[n-1] = 0
        cases, names, Ws = getTestMatrices(n)
        for i, (Z, M) in enumerate(cases):
            tau = getReducedContractionFactor(Z, M, ls, mus)
            tau_results[names[i]].append(tau)


    # Convert results to df
    tau_df = pd.DataFrame(tau_results)
    tau_df.index = range(start, end, 2)

    return tau_df

# Figure 6b
tau_df = testFixedRateProj(start=6, end=31)

Testing optimal rate on 6 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 8 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 10 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 12 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 14 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 16 nodes

# Plot figure 6b
fig, ax = plt.subplots()
tau_df.plot(ax=ax)

ax.set_xlabel("n")
ax.set_ylabel("Contraction")
ax.set_title("Fixed step size contraction with an unrestricted monotone operator")
ax.legend(['Full', 'MT', '2-Block Min SLEM', r'2-Block Min Resistance (also min $\|Z-W\|$)', '2-Block Max Fiedler'], loc="lower right")
fig.savefig("figs/fig6b.pdf", format="pdf", bbox_inches="tight")

# Figure 7 helper
def testOptimalRate(start=4, end=31, lc=2, mc=1):
    tau_results = defaultdict(list) 
    for n in range(start, end, 2):
        print(f"Testing optimal rate on {n} nodes")
        ls = np.ones(n) * lc
        ms = np.ones(n) * mc
        cases, names, Ws = getTestMatrices(n)
        for i, (Z, M) in enumerate(cases):
            _, tau = getReducedGamma(Z, M, ls, ms)
            tau_results[names[i]].append(tau)


    # Convert results to df
    tau_df = pd.DataFrame(tau_results)
    tau_df.index = range(start, end, 2)

    return tau_df

# Figure 7 helper
def plotTau(tau_df, title):
    fig, ax = plt.subplots()
    tau_df.plot(ax=ax)
    ax.set_xlabel("n")
    ax.set_ylabel("Contraction") #$\\tau$
    ax.set_title(title)
    ax.legend(['Full', 'MT', '2-Block Min SLEM', r'2-Block Min Resistance (also min $\|Z-W\|$)', '2-Block Max Fiedler'], loc="center right")
    #ax.legend(loc='center right')
    #plt.show()
    return fig

# Figure 7
tau_df = testOptimalRate(start=4, end=31)

tau_df.to_csv("figs/fig_opt_gamma_strong.csv", index=False)

Testing optimal rate on 4 nodes
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 6 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 8 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 10 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 12 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 14 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

Testing optimal rate on 16 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 18 nodes

# Figure 7
#tau_df = pd.load_csv("fig_opt_gamma_strong.csv")
fig = plotTau(tau_df, "Contraction over $n$ Lipschitz strongly monotone operators")
fig.savefig("figs/fig_opt_gamma_strong.pdf", format="pdf", bbox_inches="tight")

# Figure 7 helper
def testOptimalRateProj(start=4, end=31, lc=2, mc=1):
    tau_results = defaultdict(list) 
    for n in range(start, end, 2):
        print(f"Testing optimal rate on {n} nodes")
        ls = np.ones(n) * lc
        ls[n-1] = np.inf
        mus = np.ones(n) * mc
        mus[n-1] = 0
        cases, names, Ws = getTestMatrices(n)
        for i, (Z, M) in enumerate(cases):
            _, tau = getReducedGamma(Z, M, ls, mus)
            tau_results[names[i]].append(tau)


    # Convert results to df
    tau_df = pd.DataFrame(tau_results)
    tau_df.index = range(start, end, 2)

    return tau_df

# Figure 7b
tau_df = testOptimalRateProj(start=4, end=31)

tau_df.to_csv("figs/fig_opt_gamma_proj.csv", index=False)

Testing optimal rate on 4 nodes
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 6 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 8 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 10 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 12 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 14 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

Testing optimal rate on 16 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 18 nodes

# Figure 7b
fig = plotTau(tau_df, "Contraction with an unrestricted monotone operator")

fig.savefig("figs/fig7b.pdf", format="pdf", bbox_inches="tight")

# Figure 8 - operator ordering
Lc=2
mu=1


taus = defaultdict(list)
for n in range(4, 31, 2):
    print(f"Testing optimal rate on {n} nodes")
    Lb, Wb = getBlockMin(n, n//2)
    Mb = getMfromWCholesky(Wb)
    Lmt, Wmt = getMT(n)
    Mmt = getMfromWCholesky(Wmt)
    Lr, Wr = getRyu(n)
    Mr = getMfromWCholesky(Wr)
    Lmax, Wmax = getFull(n)
    Mmax = getMfromWCholesky(Wmax)

    cases = [(Lb, Mb), (Lmt, Mmt), (Lr, Mr), (Lmax, Mmax)]
    names = ["2-Block", "MT", "Ryu", "Full"]
    ls = np.ones(n)*Lc
    ls[n-1] = np.inf
    mus = np.ones(n)*mu
    mus[n-1] = 0
    for i, (Z, M) in enumerate(cases):
        _, tau = getReducedGamma(Z, M, ls, mus)
        taus[names[i]].append(tau)

    ls = np.ones(n)*Lc
    ls[0] = np.inf
    mus = np.ones(n)*mu
    mus[0] = 0
    for i, (Z, M) in enumerate(cases):
        _, tau = getReducedGamma(Z, M, ls, mus)        
        taus[names[i]+" reverse"].append(tau)

Testing optimal rate on 4 nodes
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 6 nodes
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 8 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 10 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 12 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 14 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
Testing optimal rate on 16 nodes

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator
No operators provided, using default LipschitzStronglyMonotoneOperator

# Figure 8 - operator ordering
df = pd.DataFrame(taus)
df.index = range(4, 31, 2)

# Plot the results
colors = ['blue', 'green', 'red', 'purple']
styles = ['b-', 'g-', 'k-', 'm-', 'y--', 'c--', 'r--', 'p--']
styles = ['b-', 'g-', 'k-', 'm-', 'b*-', 'g*-', 'k*-', 'm*-']

df.to_csv("figs/fig_function_ordering_opt.csv", index=False)
df.plot(style=styles)
plt.xlabel("Number of nodes")
plt.ylabel("Contraction")
plt.title("Operator Ordering and Splitting Performance")

plt.savefig("figs/fig_function_order_opt.pdf", format="pdf", bbox_inches="tight")
plt.show()

# Figure 9 - block size
# See 04_23_24_block_size.ipynb
# Function to generate list of all test matrices
def getTestMatricesBlocks(n):
    cases = []
    names = ['Full', 'MT']
    objs = ['Full', 'MT']
    blocksizes = ['', '']
    blocklist = ['', '']
    Lf, Wf = getFull(n)
    Mf = getMfromWCholesky(Wf)
    cases.append((Lf, Mf))
    Lmt, Wmt = getMT(n)
    Mmt = getMfromWCholesky(Wmt)
    cases.append((Lmt, Mmt))
    objstring = ['MinSLEM', 'MinResist', 'getMaxConnectivity', 'Min |Z-W|']
    for i in range(2, n//2+1):
        if n % i != 0:
            continue
        for j, obj in enumerate([getMinSLEM, getMinResist, getMaxConnectivity, getSimilar]):
            L, W = getBlockMin(n, i, objective=obj)
            M = getMfromWCholesky(W)
            cases.append((L, M))
            objs.append(objstring[j])
            blocksizes.append(i)
            blocklist.append(n//i)
            title = f"{n//i}-Block "+objstring[j]
            names.append(title)
        
    return cases, names, objs, blocksizes, blocklist

# Test function
def testMatrices(cases, names, lc=2, mu=1, gamma=0.5):
        
    taus = {}
    for i, (L, M) in enumerate(cases):
        ls = np.ones(len(L))*lc
        mus = np.ones(len(L))*mu
        
        gamma, tau = getReducedGamma(L, M, ls, mus)
        print(names[i], tau)
        taus[names[i]] = tau

    return taus

# Figure 9
results = []
for n in range(4, 25): #t < 120:
    print(n)
    cases, names, objs, blocksizes, blocks = getTestMatricesBlocks(n)
    taus = testMatrices(cases, names)
    # Each key in taus is a column, the row is the value of n
    for i in range(len(names)):
        k = names[i]
        v = taus[k]
        results.append({"n": n, "name": k, "obj": objs[i], "blocksize": blocksizes[i], "blocks":blocks[i], "tau": v})

df = pd.DataFrame(results)

# Save to csv
df.to_csv("figs/taus_plus.csv", index=False)
df.set_index('n', inplace=True)

4

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.4275726424387023
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.8002281453197694
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.5571563320386254
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinResist 0.4285714308439043
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.47232795740698197
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block Min |Z-W| 0.42857142974429335
5
No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.44769222096754585
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.861371784044431
6

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.45925820423443237
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.8973983347682633
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinSLEM 0.7135742120391384
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinResist 0.7182631442291072
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block getMaxConnectivity 0.7125769076136125
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block Min |Z-W| 0.7146237387666267
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.5571562848351927
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinResist 0.4285714345044546
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.488471357684012
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block Min |Z-W| 0.4285714294945974
7
No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.46675734145403897
No operators provided, using default LipschitzStronglyMonotoneOperator

MT 0.921028546326316
8

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.47200715470169324
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9373167138798217
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block MinSLEM 0.8309037714895217
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block MinResist 0.8108953220221728
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block getMaxConnectivity 0.8002349980331156
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block Min |Z-W| 0.9414775041859474
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.5571562864004128
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinResist 0.4285714324225226
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.49712977590980384
No operators provided, using default LipschitzStronglyMonotoneOperator

2-Block Min |Z-W| 0.4285717289200378
9

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.47588407918368214
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9490542144439524
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinSLEM 0.7139323261601181
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinResist 0.7182631341824306
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block getMaxConnectivity 0.7138583593905867
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block Min |Z-W| 0.7154479936264707
10

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.4788623869914284
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9578086304484562
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block MinSLEM 0.8909710630474715
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block MinResist 0.867517704847014
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block getMaxConnectivity 0.8694665060801878
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block Min |Z-W| 0.8841704218835689
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.5571562868119515
No operators provided, using default LipschitzStronglyMonotoneOperator

2-Block MinResist 0.4285714375496707
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.5007530033260064
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block Min |Z-W| 0.4285714341244909
11
No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.48122091775576264
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9645090300427932
12

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.48313426706054324
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9697458591679796
No operators provided, using default LipschitzStronglyMonotoneOperator
6-Block MinSLEM 0.9274863710223136
No operators provided, using default LipschitzStronglyMonotoneOperator
6-Block MinResist 0.9043223463150312
No operators provided, using default LipschitzStronglyMonotoneOperator
6-Block getMaxConnectivity 0.914435584101685
No operators provided, using default LipschitzStronglyMonotoneOperator

df[df["obj"].isin(["MinResist", "Full", "MT"])].groupby(['obj', 'blocks'])['tau'].groups.keys()

dict_keys([('Full', ''), ('MT', ''), ('MinResist', 2), ('MinResist', 3), ('MinResist', 4), ('MinResist', 5), ('MinResist', 6), ('MinResist', 7), ('MinResist', 8), ('MinResist', 9), ('MinResist', 10), ('MinResist', 11), ('MinResist', 12)])

# Figure 9
# Load taus.csv
#df = pd.read_csv("taus.csv")
#df.loc[df['name'].isin(['Full', 'MT']), ['blocks', 'blocksize']] = ''

df[df["obj"].isin(["MinResist", "Full", "MT"])].groupby(['obj', 'blocks'])['tau'].plot(marker='o', linestyle='-', figsize=(10, 6))

keys = df[df["obj"].isin(["MinResist", "Full", "MT"])].groupby(['obj', 'blocks'])['tau'].groups.keys()
keys = ['Full', 'MT', '2-Block', '3-Block', '4-Block', '5-Block', '6-Block', '7-Block', '8-Block', '9-Block', '10-Block', '11-Block', '12-Block']
plt.legend(keys, bbox_to_anchor=(1, 1))
# Make x axis integer
plt.xticks(np.arange(4, 25, 2))
plt.xlabel(r"$n$")
plt.title("Min Resistance Contraction over $n$ Lipschitz Strongly Monotone Operators by Block Number")
plt.savefig("figs/fig_block_count.pdf", format="pdf", bbox_inches="tight")

# Figure 9 - Test function
def testMatricesProj(n, cases, names, lc=2, mu=1):
    ls = np.ones(n)*lc
    mus = np.ones(n)*mu
    ls[n-1] = np.inf
    mus[n-1] = 0
    taus = {}
    for i, (Z, M) in enumerate(cases):
        gamma, tau = getReducedGamma(Z, M, ls, mus)
        print(names[i], tau)
        taus[names[i]] = tau

    return taus

# Figure 9
projresults = []
for n in range(4,26): #t < 120:
    print(n)
    cases, names, objs, blocksizes, blocks = getTestMatricesBlocks(n)
    taus = testMatricesProj(n, cases, names)
    # Each key in taus is a column, the row is the value of n
    for i in range(len(names)):
        k = names[i]
        v = taus[k]
        
        projresults.append({"n": n, "name": k, "obj": objs[i], "blocksize": blocksizes[i], "blocks":blocks[i], "tau": v})

projdf = pd.DataFrame(projresults)
projdf.to_csv("figs/taus_proj.csv", index=False)
projdf.set_index('n', inplace=True)

4

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.7521521274547772
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9177841369745814
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.8191387832657023
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinResist 0.7870338922103205
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.8005884142677305
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block Min |Z-W| 0.7870338929300077
5
No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.8044792127888197
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9475773165769197
6

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.838558740668551
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.963679040723567
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinSLEM 0.9045196638648727
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinResist 0.8989263093339976
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block getMaxConnectivity 0.9031665201664806
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block Min |Z-W| 0.905811564931693
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.8708680920365564
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinResist 0.8616700832103537
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.8648381115704538
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block Min |Z-W| 0.8616700813867951
7
No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.8625324131656744
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9733944272799456
8

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.8803148080938116
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9796953548701708
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block MinSLEM 0.947737890434155
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block MinResist 0.9366150047051096
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block getMaxConnectivity 0.9409068724807719
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block Min |Z-W| 0.9733653867324151
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.899665656605118
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinResist 0.897358435818429
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.897061358320355
No operators provided, using default LipschitzStronglyMonotoneOperator

2-Block Min |Z-W| 0.8973584363622054
9

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.8940287042166419
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9840061423604513
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinSLEM 0.9310037419763251
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinResist 0.9226238048816924
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block getMaxConnectivity 0.9308683371509562
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block Min |Z-W| 0.9329235479910636
10

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.904926183461662
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9870820594200824
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block MinSLEM 0.9690123429432841
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block MinResist 0.9580312775913491
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block getMaxConnectivity 0.9629973163283422
No operators provided, using default LipschitzStronglyMonotoneOperator
5-Block Min |Z-W| 0.9629875076136046
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block MinSLEM 0.9180267565103261
No operators provided, using default LipschitzStronglyMonotoneOperator

2-Block MinResist 0.9183673760071706
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block getMaxConnectivity 0.916691539095593
No operators provided, using default LipschitzStronglyMonotoneOperator
2-Block Min |Z-W| 0.9183673754845947
11
No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.9137934450814554
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9893521326451261
12

No operators provided, using default LipschitzStronglyMonotoneOperator
Full 0.9211491280506094
No operators provided, using default LipschitzStronglyMonotoneOperator
MT 0.9910743537751981
No operators provided, using default LipschitzStronglyMonotoneOperator
6-Block MinSLEM 0.9794372274110801
No operators provided, using default LipschitzStronglyMonotoneOperator
6-Block MinResist 0.9701555310266454
No operators provided, using default LipschitzStronglyMonotoneOperator
6-Block getMaxConnectivity 0.9756973937657129
No operators provided, using default LipschitzStronglyMonotoneOperator

6-Block Min |Z-W| 0.9901053735796335
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block MinSLEM 0.9594826252987537
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block MinResist 0.9489711986480376
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block getMaxConnectivity 0.9564324332772989
No operators provided, using default LipschitzStronglyMonotoneOperator
4-Block Min |Z-W| 0.9875641251533125
No operators provided, using default LipschitzStronglyMonotoneOperator
3-Block MinSLEM 0.946111764232151
No operators provided, using default LipschitzStronglyMonotoneOperator

# Figure 9 Class 2 

#projdf = pd.read_csv("figs/taus_proj.csv")
projdf[projdf["obj"].isin(["MinResist", "Full", "MT"])].groupby(['obj', 'blocks'])['tau'].plot(marker='o', linestyle='-', figsize=(10, 6))

keys = ['Full', 'MT', '2-Block', '3-Block', '4-Block', '5-Block', '6-Block', '7-Block', '8-Block', '9-Block', '10-Block', '11-Block', '12-Block'] # If using df directly
plt.legend(keys, bbox_to_anchor=(1, 1))
plt.title("Min Resistance Contraction over $n-1$ Lipschitz Strongly Monotone Operators")
plt.savefig("figs/fig_block_count_proj.pdf", format="pdf", bbox_inches="tight")

# Figure 10 - pareto curves
# 04_21_24_full_pareto.ipynb

from matplotlib.colors import CenteredNorm
def getTaus(Z, W, l=None, mu=None):

    M = getMfromWCholesky(W)
    h = getReducedContractionFactor(Z, M, l, mu, gamma=0.5)
    m = getReducedGamma(Z, M, l, mu)[1]
    return h, m

def getContractTime(n, l, mu, objective=None, tol=0.5):
    Z_fixed = {}

    t = np.ones(n)
    ll = np.ones((n,n))
    results = {}
    for i in range(n//2):
        print(i)
        start = 0
        for k in range(2):
            for j in range(start, i + k*(n//2)):
                Z_fixed[(i + k*(n//2),j)] = 0
            Zs, Ws = objective(n, fixed_Z=Z_fixed)
            h, m = getTaus(Zs, Ws, l, mu)
            tt, _, _ = getCycleTime(t, ll, Zs, Ws, itrs=2)
            # r is number of entries in Z_fixed
            r = len(Z_fixed)
            tif = np.ceil(np.log(tol)/np.log(h)) # Total iterations fixed
            tio = np.ceil(np.log(tol)/np.log(m)) # Total iterations opt step
            ttf = tif*tt
            tto = tio*tt
            results[r] = {'Fixed':h , 'M Opt': m, 'Time': tt, 'Removed Edges': r, 'Fixed Total Iterations': tif, 'Opt Step Total Iterations': tio, 'Fixed Total Time': ttf, 'Opt Step Total Time': tto}
            print(results[r])
            start = n//2

    counter = 0
    for i in range(n//2 - 2):
        for j in range(i+1):
            r = n//2 + j
            c = n//2 - 1 - i + j
            Z_fixed[(r, c)] = 0
            # print(r,c)
            counter += 1
            if counter % 10 == 0:
                Ls, Ws = objective(n, fixed_Z=Z_fixed)
                h, m = getTaus(Ls, Ws, l, mu)
                tt = 4
                # r is number of entries in Z_fixed
                rem = len(Z_fixed)    
                tif = np.ceil(np.log(tol)/np.log(h)) # Total iterations fixed
                tio = np.ceil(np.log(tol)/np.log(m)) # Total iterations opt step
                ttf = tif*tt
                tto = tio*tt
                results[rem] = {'Fixed':h , 'M Opt': m, 'Time': tt, 'Removed Edges': rem, 'Fixed Total Iterations': tif, 'Opt Step Total Iterations': tio, 'Fixed Total Time': ttf, 'Opt Step Total Time': tto}
                print(results[rem])
    # Create dataframe of taus and times
    df = pd.DataFrame(results)

    return df

# Figure 10
n = 24
l = np.ones(n)*2
l[n-1] = np.inf
mu = np.ones(n)
mu[n-1]= 0

df = getContractTime(n, l, mu, objective=getMaxConnectivity)
df2 = df.T

0

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator

{'Fixed': 0.9798415277932965, 'M Opt': 0.9610519693854286, 'Time': 47.99999999999555, 'Removed Edges': 0, 'Fixed Total Iterations': 35.0, 'Opt Step Total Iterations': 18.0, 'Fixed Total Time': 1679.9999999998442, 'Opt Step Total Time': 863.99999999992}

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator

{'Fixed': 0.9798415277932965, 'M Opt': 0.9610519693854286, 'Time': 47.99999999999555, 'Removed Edges': 0, 'Fixed Total Iterations': 35.0, 'Opt Step Total Iterations': 18.0, 'Fixed Total Time': 1679.9999999998442, 'Opt Step Total Time': 863.99999999992}
1

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator

{'Fixed': 0.9802934213873657, 'M Opt': 0.9611344883825657, 'Time': 46.00000000002653, 'Removed Edges': 1, 'Fixed Total Iterations': 35.0, 'Opt Step Total Iterations': 18.0, 'Fixed Total Time': 1610.0000000009286, 'Opt Step Total Time': 828.0000000004776}

df.to_csv('figs/n24.csv', index=False)
# df = pd.read_csv("figs/n24.csv")

# Plot total time vs removed edges
plt.plot(df.iloc[3], df.iloc[6], df.iloc[3], df.iloc[7])
# Subtitle
plt.suptitle('Time Units to 0.5 Contraction vs Removed Edges, $n=$'+str(n))


plt.legend( ['Fixed Step Size', 'Optimal Step Size'])
plt.xlabel('Removed Edges')
plt.ylabel('Time Units')
plt.savefig("figs/fig_optimal_sparsity.pdf", format="pdf", bbox_inches="tight")
plt.show()

# Figure 10
# Plot the top three rows of df as a function of time with points annotated
plt.plot(df.iloc[2], df.iloc[0], marker='<')
plt.plot(df.iloc[2], df.iloc[1], marker='<')

plt.title('Contraction vs Time as Edges are Removed')
# Add legend
plt.legend( ['Fixed Step Size', 'Optimal Step Size'])
# Annotate points with number of removed edges
lx, ly = (0,0)
right = 0
up = 10
for j in range(len(df.columns)):
        print(j,(df.iloc[2,j], df.iloc[1,j]))
        if df.iloc[2,j] < 4:
                right = 10
                up = 1
        if df.iloc[2,j] == 4.0:
                up = 0
        if abs(lx - df.iloc[2,j]) > 1 or abs(ly - df.iloc[1,j]) > 0.001:
                plt.annotate(df.columns[j], (df.iloc[2,j], df.iloc[0,j]), textcoords="offset points", xytext=(right,up), ha='center', fontsize=8)
                plt.annotate(df.columns[j], (df.iloc[2,j], df.iloc[1,j]), textcoords="offset points", xytext=(right,up), ha='center', fontsize=8)
                lx = df.iloc[2,j]
                ly = df.iloc[1,j]
        
# Annotate with points

plt.xlabel('Average Iteration Time')
plt.ylabel(r'Contraction Rate ($\tau$)')

plt.savefig("figs/fig_pareto.pdf", format="pdf", bbox_inches="tight")
plt.show()

0 (47.99999999999555, 0.9610519693854286)
1 (46.00000000002653, 0.9611344883825657)
2 (44.00000000001398, 0.9612008911749754)
3 (41.99999999999353, 0.9612841961043996)
4 (40.00000000033539, 0.9613404324143697)
5 (38.00000000019667, 0.9614511688616295)
6 (36.00000000005644, 0.9615160720131097)
7 (34.00000000000092, 0.9616399278912013)
8 (32.00000000018983, 0.9616502219370228)
9 (29.999999999934413, 0.961814464869611)
10 (28.000000000015866, 0.9618911233993571)
11 (26.000000000192202, 0.9619696358597921)
12 (23.9999999999246, 0.9620812351430117)
13 (22.000000000001247, 0.9622244218627922)
14 (19.99999999987995, 0.962304209896036)
15 (17.999999999943523, 0.9624258907775105)
16 (15.999999999994685, 0.9624910820032222)
17 (14.000000000000364, 0.9627105025485819)
18 (11.9999999999999, 0.9627679103016183)
19 (10.000000000165418, 0.9630429010126391)
20 (7.999999999999933, 0.9630413776592525)
21 (3.999999999999991, 0.9642038633941726)
22 (3.999999999999991, 0.9642023236908612)
23 (4.0, 0.9652843389466005)
24 (4.0, 0.966900228979608)
25 (4.0, 0.9686471381342026)
26 (4.0, 0.9699045174130262)
27 (4.0, 0.9724134559398137)

# Figure 11 - edge deletion ordering
# 04_21_24_full_pareto.ipynb

def cplot(M, c, title=''):
    vmin = 1
    vmax = np.max(ordermat)
    cm = plt.cm.viridis #coolwarm
    fig, ax = plt.subplots(1, 1, figsize=(12, 6))
    x = np.arange(n)
    y = -np.arange(n)
    pc = ax.pcolormesh(x, y, M, cmap=cm , vmin=vmin, vmax=vmax)
    ax.set_title(title) #'Edge Removal Ordering'
    ax.axis('off')
    ax.set_aspect('equal')

    # Add colorbar
    plt.tight_layout()
    
    cbar = fig.colorbar(pc) 
    cbar.set_ticks(np.arange(1, c, 4)) 
    return fig

# Display edge removal ordering
def order_graph(n):
    order_matrix = np.zeros((n,n)) 
    counter = 1
    for i in range(n//2):
        start = 0
        for k in range(2):
            for j in range(start, i + k*(n//2)):
                r = i + k*(n//2)
                c = j
                order_matrix[(r,j)] = counter
                order_matrix[(c,r)] = counter
            
            counter += 1
            start = n//2

    ccounter = 0
    for i in range(n//2 - 2):
        for j in range(i+1):
            r = n//2 + j
            c = n//2 - 1 - i + j
            order_matrix[(r, c)] = counter
            order_matrix[(c, r)] = counter
            ccounter += 1
            if ccounter % 10 == 0:
                counter += 1
            
    return order_matrix, counter
ordermat, c = order_graph(24)
ordermat = np.ma.masked_array(ordermat, ordermat < 1)

fig = cplot(ordermat, c, 'Edge Removal Ordering')
fig.savefig("figs/fig_edge_removal.pdf")

# Figure 12
block_results = []
min_results = []
for n in range(6,10):
    for i in range(40):
        t = np.random.rand(n)*10+1

        # Communication time
        l = np.random.rand(n, n)*1.5 + 0.5
        l = np.tril(l, -1) + np.tril(l, -1).T
        Zb, Wb = getBlockMin(n, n//2, objective=getMinResist)
        c, tau, cc, tt = getMetrics(Zb, Wb, t, l,contraction_target=0.01)
        block_results.append((n, i, c, tau, cc, tt))
        Z, W = getMinCycle(n, objective=getMinResist, t=t, l=l, c=0.01, r=n, minW=n//2)
        if Z is not None:
            c, tau, cc, ttc = getMetrics(Z, W, t, l,contraction_target=0.01)
        else:
            c, tau, cc, ttc = (None, None, None, None)
        min_results.append((n, i, c, tau, cc, ttc))
        print(n,i,tt,ttc)

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
6 0 557.3531756387047 557.3531756387047

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
6 1 618.6495759991482 522.5575943515803

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
6 2 634.3218614903401 634.3218614903401

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
6 3 641.4972132053769 641.4972132053769

No operators provided, using default LipschitzStronglyMonotoneOperator

No operators provided, using default LipschitzStronglyMonotoneOperator
6 4 698.1699747774383 698.1699747774383

# Figure 12
block_df = pd.DataFrame(block_results)
min_df = pd.DataFrame(min_results)

block_df.to_csv('figs/block_df2.csv')
min_df.to_csv('figs/min_cycle_min_edge.csv')

# Paired density chart for block vs min cycle
# Figure 12
# block_df = pd.read_csv('figs/block_df2.csv')
# min_df = pd.read_csv('figs/min_cycle_min_edge.csv')
fig, axs = plt.subplots(2, 2, sharey=True)
n = 6
key = 5 #'5' #use string if loading from csv
start = 0
for i in range(2):
    for j in range(2):
        end = 40 + start
        axs[i,j] = block_df[key][start:end].plot.kde(ax=axs[i,j], color='blue')
        min_df[key][start:end].plot.kde(ax=axs[i,j], color='red')
        start += 40
        axs[i,j].set_title(r"$n=$"+str(n))
        n += 1

fig.legend([r'$d$-Block Design', 'Con. Min Iter Time'], loc=(.65,.8))

for ax in axs.flat:
    ax.set(xlabel='Time')
for ax in fig.get_axes():
    ax.label_outer()
fig.savefig("figs/fig_constrained_mincycle_histogram.pdf", format="pdf", bbox_inches="tight")

Optimized Resolvent Splitting Examples¶